from os.path import splitext
from qgis.core import (
    QgsRasterLayer, QgsProject, QgsCoordinateTransformContext, QgsApplication, QgsVectorLayer, QgsCoordinateReferenceSystem, QgsCoordinateTransform,
    QgsProcessingFeedback, QgsProcessingUtils,
    QgsVectorFileWriter,QgsProcessingRegistry
)
from qgis.analysis import (QgsRasterCalculator, QgsRasterCalculatorEntry)
import os
import processing
from processing.core.Processing import Processing
from processing.algs.gdal.GdalAlgorithmProvider import GdalAlgorithmProvider
from osgeo import gdal,ogr
# 初始化QGIS应用（若未初始化）




def filter_rgb(layer, a1, a2, a3, out_path=None, out_name=None):
    """见原文档"""
    if isinstance(layer, str):           # 允许直接传路径
        layer = QgsRasterLayer(layer, 'temp')

    if not layer.isValid():
        raise ValueError("图层无效")

    if out_path is None:
        base, _ = splitext(layer.source())
        out_path = base + "_filtered.tif"

    entries = []
    for b in (1, 2, 3):
        e = QgsRasterCalculatorEntry()
        e.ref = f"{layer.name()}@{b}"
        e.raster = layer
        e.bandNumber = b
        entries.append(e)

    # 修正格式化语法
    expr = (
        f'("{layer.name()}@1" = {a1}) * '
        f'("{layer.name()}@2" = {a2}) * '
        f'("{layer.name()}@3" = {a3})'
    )
    print(expr)

    calc = QgsRasterCalculator(
        expr, out_path, "GTiff",
        layer.extent(), layer.width(), layer.height(),
        entries, QgsCoordinateTransformContext()
    )

    if calc.processCalculation() != 0:
        raise RuntimeError("栅格计算器执行失败")

    out_layer = QgsRasterLayer(out_path, out_name or f"{layer.name()}_filtered")
    if not out_layer.isValid():
        raise RuntimeError("无法加载输出栅格")
    QgsProject.instance().addMapLayer(out_layer)
    return out_layer

def safe_polygonize(input_raster, output_shp,band = 1):
    try:
        # 1. 矢量化
        raster_ds = gdal.Open(input_raster)
        if raster_ds is None:
            raise ValueError("无法打开栅格文件")
        driver = ogr.GetDriverByName("ESRI Shapefile")
        vector_ds = driver.CreateDataSource(output_shp)
        layer = vector_ds.CreateLayer("polygon", geom_type=ogr.wkbPolygon)
        field_defn = ogr.FieldDefn("value", ogr.OFTInteger)
        layer.CreateField(field_defn)
        band = raster_ds.GetRasterBand(band)
        gdal.Polygonize(band, None, layer, 0, [], callback=None)
        vector_ds = None
        raster_ds = None
        print(f"结果已保存至: {output_shp}")
    except Exception as e:
        print(f"矢量化失败: {str(e)}")


# ---------- 调用 ----------
safe_polygonize(r"I:\大论文\示范区栅格\高程.tif",r"I:\大论文\示范区栅格\高程.shp")
#filter_rgb(r"I:\大论文\示范区栅格\高程.tif", 111, 222, 333)